#!/usr/bin/env python
# -*- coding: utf-8 -*-
from gensim.summarization import bm25
import joblib
from config import Config


class BM25(object):
    """
    基于bm25文本匹配
    """

    def __init__(self, status=False):
        self.filename = Config.BM25_MODEL
        self.word_name = Config.BM25_MODEL_WORDS
        if not status:
            self.bm25_Model = joblib.load(filename=self.filename)
            self.bm25_data = joblib.load(filename=self.word_name)
            self.average_idf = sum(map(lambda k: float(self.bm25_Model.idf[k]), self.bm25_Model.idf.keys())) / len(
                self.bm25_Model.idf.keys())

    def train(self, corpus: list):
        """
        保存待匹配的语料
        :param corpus:语料
        """
        bm25_Model = bm25.BM25([list(element) for element in corpus])
        joblib.dump(bm25_Model, self.filename)
        joblib.dump([element for element in corpus], self.word_name)

    def predict(self, text: str) -> list:
        """
        匹配最相近的文本
        :param text: 待匹配的文本
        :return:排序后匹配的内容
        """
        scores = self.bm25_Model.get_scores([element for element in text], self.average_idf)
        data = dict(zip(self.bm25_data, scores))
        data = sorted(data.items(), key=lambda item: item[1], reverse=True)
        data = [element for element in data if element[1] > 0]
        return data


if __name__ == '__main__':
    # BM25(status=True).train(list(set([line.strip("\n") for line in open("../../data/BM25/del.txt",encoding="utf8")])))
    print(BM25().predict("肚子疼"))